package io.github.kuafucv.zero.log.io;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.concurrent.BlockingDeque;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.LinkedBlockingDeque;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.ReentrantLock;

public class DiscardingRollingOutput {

    private ByteArrayOutputStream currentBucket;

    private final ReentrantLock currentBucketLock = new ReentrantLock();

    private final BlockingDeque<ByteArrayOutputStream> filledBuckets;

    private final ConcurrentLinkedQueue<ByteArrayOutputStream> recycledBucketPool;

    private final long maxBucketSizeInBytes;

    private final AtomicInteger discardedBucketCount = new AtomicInteger();

    /**
     * @param maxBucketSizeInBytes 每个bucket的最大字节大小
     * @param maxBucketCount 最大桶数
     */
    public DiscardingRollingOutput(int maxBucketSizeInBytes, int maxBucketCount) {
        if (maxBucketCount < 2) {
            throw new IllegalArgumentException("'maxBucketCount' must be >1");
        }

        this.maxBucketSizeInBytes = maxBucketSizeInBytes;
        this.filledBuckets = new LinkedBlockingDeque<>(maxBucketCount);
        this.recycledBucketPool = new ConcurrentLinkedQueue<>();
        this.currentBucket = newBucket();
    }


    public void write(byte[] b) throws IOException {
        currentBucketLock.lock();
        try {
            if(currentBucket.size() == 0){
                currentBucket.write("[".getBytes(StandardCharsets.UTF_8));
            }else {
                currentBucket.write(",".getBytes(StandardCharsets.UTF_8));
            }
            currentBucket.write(b);
            rollCurrentBucketIfNeeded();
        } finally {
            currentBucketLock.unlock();
        }
    }

    /**
     * 如果已达到大小阈值，则滚动当前存储桶。
     */
    private void rollCurrentBucketIfNeeded() throws IOException {
        if (currentBucket.size() < maxBucketSizeInBytes) {
            return;
        }
        rollCurrentBucket();
    }

    /**
     * 只要桶大小不为空，强制滚动当前存储桶。
     */
    public void rollCurrentBucketIfNotEmpty() throws IOException{
        if (currentBucket.size() == 0) {
            return;
        }
        rollCurrentBucket();
    }

    /**
     * 将当前活动的存储桶移动到已填充的存储桶列表中，并定义一个新的存储桶。
     */
    public void rollCurrentBucket() throws IOException{
        currentBucketLock.lock();
        try {
            currentBucket.write("]".getBytes(StandardCharsets.UTF_8));
            boolean offered = filledBuckets.offer(currentBucket);
            if (offered) {
                onBucketRoll(currentBucket);
            } else {
                onBucketDiscard(currentBucket);
                discardedBucketCount.incrementAndGet();
            }
            currentBucket = newBucket();
        } finally {
            currentBucketLock.unlock();
        }
    }

    protected ByteArrayOutputStream newBucket() {
        ByteArrayOutputStream bucket = recycledBucketPool.poll();
        if (bucket == null) {
            bucket = new ByteArrayOutputStream();
        }
        return bucket;
    }

    /**
     * 将给定的bucket返回到空闲bucket池中。
     */
    public void recycleBucket(ByteArrayOutputStream bucket) {
        bucket.reset();
        recycledBucketPool.offer(bucket);
    }

    /**
     * 归还装满的桶
     */
    public BlockingDeque<ByteArrayOutputStream> getFilledBuckets() {
        return filledBuckets;
    }

    /**
     * 丢弃的桶数。以监控为导向的指标。
     */
    public int getDiscardedBucketCount() {
        return discardedBucketCount.get();
    }

    public long getCurrentOutputStreamSize() {
        long sizeInBytes = 0;
        for (ByteArrayOutputStream bucket : filledBuckets) {
            sizeInBytes += bucket.size();
        }
        sizeInBytes += currentBucket.size();
        return sizeInBytes;
    }

    /**
     * 关闭所有底层桶（当前桶、装满的桶和回收桶池中的桶）。
     */
    public void close() {
    }

    /**
     * 专为扩展而设计。
     */
    protected void onBucketDiscard(ByteArrayOutputStream discardedBucket) {
    }

    /**
     * 滚动的桶。专为扩展而设计。
     */
    protected void onBucketRoll(ByteArrayOutputStream rolledBucket) {
    }

    @Override
    public String toString() {
        return "DiscardingRollingOutputStream{" +
                "currentBucket.bytesWritten=" + currentBucket.size() +
                ", filledBuckets.size=" + filledBuckets.size() +
                ", discardedBucketCount=" + discardedBucketCount +
                ", recycledBucketPool.size=" + recycledBucketPool.size() +
                '}';
    }

}
